import os
import torch
from torch.utils.data import DataLoader
from torch.optim import Adam
from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts
from tqdm import tqdm

from src.config.config import Config
from src.data.dataset import DigitsDataset
from src.models.resnet import DigitsResnet18, DigitsResnet50
from src.models.mobilenet import DigitsMobilenet
from src.models.loss import LabelSmoothEntropy
from src.utils.metrics import accuracy

class Trainer:
    def __init__(self, config, model_name='resnet18', val=True):
        self.config = config
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        
        # 初始化数据加载器
        self.train_set = DigitsDataset(config.data_dir, mode='train')
        self.train_loader = DataLoader(
            self.train_set, 
            batch_size=config.batch_size,
            shuffle=True,
            num_workers=8,
            pin_memory=True,
            drop_last=True,
            collate_fn=self.train_set.collect_fn
        )

        if val:
            self.val_loader = DataLoader(
                DigitsDataset(config.data_dir, mode='val', aug=False),
                batch_size=config.batch_size,
                num_workers=8,
                pin_memory=True,
                drop_last=False
            )
        else:
            self.val_loader = None

        # 初始化模型
        if model_name == 'resnet18':
            self.model = DigitsResnet18(config.class_num)
        elif model_name == 'resnet50':
            self.model = DigitsResnet50(config.class_num)
        elif model_name == 'mobilenet':
            self.model = DigitsMobilenet(config.class_num)
        else:
            raise ValueError(f"Unknown model name: {model_name}")
            
        self.model = self.model.to(self.device)

        # 初始化损失函数和优化器
        self.criterion = LabelSmoothEntropy(smooth=config.smooth).to(self.device)
        self.optimizer = Adam(self.model.parameters(), lr=config.lr)
        self.lr_scheduler = CosineAnnealingWarmRestarts(
            self.optimizer, 
            T_0=10, 
            T_mult=2, 
            eta_min=1e-6
        )

        self.best_acc = 0
        
        # 加载预训练模型
        if config.pretrained is not None and os.path.exists(config.pretrained):
            self.load_model(config.pretrained)
            if self.val_loader is not None:
                acc = self.eval()
                self.best_acc = acc
                print(f'Load model from {config.pretrained}, Eval Acc: {acc*100:.2f}%')

    def train(self):
        for epoch in range(self.config.start_epoch, self.config.epoches):
            train_acc = self.train_epoch(epoch)
            
            if (epoch + 1) % self.config.eval_interval == 0:
                print('Start Evaluation')
                if self.val_loader is not None:
                    val_acc = self.eval()
                    
                    if val_acc > self.best_acc:
                        os.makedirs(self.config.checkpoints, exist_ok=True)
                        save_path = os.path.join(
                            self.config.checkpoints,
                            f'epoch-{epoch+1}-acc-{val_acc*100:.2f}.pth'
                        )
                        self.save_model(save_path)
                        print(f'Model saved to {save_path}')
                        self.best_acc = val_acc

    def train_epoch(self, epoch):
        self.model.train()
        total_loss = 0
        total_acc = 0
        
        tbar = tqdm(self.train_loader, desc=f'Epoch {epoch+1}/{self.config.epoches}')
        for i, (img, label) in enumerate(tbar):
            img = img.to(self.device)
            label = label.to(self.device)
            
            self.optimizer.zero_grad()
            pred = self.model(img)
            
            # 计算损失
            loss = sum(self.criterion(p, label[:, j]) for j, p in enumerate(pred))
            total_loss += loss.item()
            
            # 反向传播
            loss.backward()
            self.optimizer.step()
            
            # 计算准确率
            acc = accuracy(pred, label)
            total_acc += acc
            
            if (i + 1) % self.config.print_interval == 0:
                self.lr_scheduler.step()
                
            tbar.set_postfix({
                'loss': f'{total_loss/(i+1):.3f}',
                'acc': f'{total_acc/(i+1)*100:.2f}%'
            })
            
        return total_acc / len(self.train_loader)

    def eval(self):
        self.model.eval()
        total_acc = 0
        
        with torch.no_grad():
            tbar = tqdm(self.val_loader, desc='Validation')
            for i, (img, label) in enumerate(tbar):
                img = img.to(self.device)
                label = label.to(self.device)
                
                pred = self.model(img)
                acc = accuracy(pred, label)
                total_acc += acc
                
                tbar.set_postfix({'acc': f'{total_acc/(i+1)*100:.2f}%'})
                
        return total_acc / len(self.val_loader)

    def save_model(self, save_path):
        torch.save({
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'scheduler_state_dict': self.lr_scheduler.state_dict(),
            'best_acc': self.best_acc,
            'epoch': self.config.start_epoch
        }, save_path)

    def load_model(self, load_path):
        checkpoint = torch.load(load_path)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        self.lr_scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
        self.best_acc = checkpoint['best_acc']
        self.config.start_epoch = checkpoint['epoch']

if __name__ == '__main__':
    config = Config()
    trainer = Trainer(config, model_name='resnet18')
    trainer.train() 